//
// Created by daybeha on 24-1-24.
//
// github :https://github.com/FeiYull/TensorRT-Alpha

#include <filesystem>
#include"deepNetwork/include/utils/yolo.h"
#include"yolov8_seg.h"

using namespace std;

void setParameters(yolo_utils::InitParameter& initParameters)
{
    initParameters.class_names = yolo_utils::dataSets::coco80;
//    initParameters.class_names = yolo_utils::dataSets::voc20;
    initParameters.num_class = 80; // for coco
//    initParameters.num_class = 20; // for voc2012
    initParameters.batch_size = 1;  // 8;
//    initParameters.src_w = 1241;
//    initParameters.src_h = 376;
    initParameters.src_w = 1280;
    initParameters.src_h = 560;
    initParameters.dst_h = 640;
    initParameters.dst_w = 640;
    initParameters.input_output_names = { "images",  "output0" };
    initParameters.conf_thresh = 0.25f;
    initParameters.iou_thresh = 0.7f;

    initParameters.is_show = true;
    initParameters.is_save = false;
}

void task(YOLOv8Seg& yolo, const yolo_utils::InitParameter& param, std::vector<cv::Mat>& imgsBatch, const int& delayTime)
{
    yolo.copy(imgsBatch);
    yolo_utils::DeviceTimer d_t1; yolo.preprocess(imgsBatch);  float t1 = d_t1.getUsedTime();
    yolo_utils::DeviceTimer d_t2; yolo.infer();				  float t2 = d_t2.getUsedTime();
    yolo_utils::DeviceTimer d_t3; yolo.postprocess(imgsBatch); float t3 = d_t3.getUsedTime();
    float avg_times[3] = { t1 / param.batch_size, t2 / param.batch_size, t3 / param.batch_size };
    sample::gLogInfo << "preprocess time = " << avg_times[0] << "; "
                     "infer time = " << avg_times[1] << "; "
                   "postprocess time = " << avg_times[2] << std::endl;
//    yolo.showAndSave(param.class_names, delayTime, imgsBatch);
    cv::Mat mask = cv::Mat(param.src_h, param.src_w, CV_8UC1, cv::Scalar(255)); // 标记点的图像
    yolo.setmask(1, imgsBatch[0], mask);  // 去除动态物体
    yolo.reset();
}

void LoadKittiImages(const string &strPathToSequence,
                     vector<string> &vstrImageLeft, vector<string> &vstrImageRight,
                     vector<double> &vTimestamps)
{
    ifstream fTimes;
    string strPathTimeFile = strPathToSequence + "/times.txt";
    fTimes.open(strPathTimeFile.c_str());
    while(!fTimes.eof())
    {
        string s;
        getline(fTimes,s);
        if(!s.empty())
        {
            stringstream ss;
            ss << s;
            double t;
            ss >> t;
            vTimestamps.push_back(t);
        }
    }

    string strPrefixLeft = strPathToSequence + "/image_0/";
    string strPrefixRight = strPathToSequence + "/image_1/";

    const int nTimes = vTimestamps.size();
    vstrImageLeft.resize(nTimes);
    vstrImageRight.resize(nTimes);

    for(int i=0; i<nTimes; i++)
    {
        stringstream ss;
        ss << setfill('0') << setw(6) << i;
        vstrImageLeft[i] = strPrefixLeft + ss.str() + ".png";
        vstrImageRight[i] = strPrefixRight + ss.str() + ".png";
    }
}

void LoadUrbanIamges(const string &strPathToSequence, vector<string> &vstrImageLeft)
{
    string strPrefixLeft = strPathToSequence + "/";
    for (const auto& entry : std::filesystem::directory_iterator(strPrefixLeft)) {
      if (entry.is_regular_file() && entry.path().extension() == ".png") {
        vstrImageLeft.push_back(entry.path().string());
//        cv::Mat image = cv::imread(entry.path().string());
//        if (image.empty()) {
//          std::cout << "Error: Unable to read image " << entry.path().string() << std::endl;
//        } else {
//          std::cout << "Reading image " << entry.path().string() << std::endl;
////          cv::imshow("Image", image);
////          cv::waitKey(0);
//        }
      }
    }
}

int main(int argc, char** argv)
{
    yolo_utils::InitParameter param;
    setParameters(param);
    std::string model_path = "/home/daybeha/Documents/XYQ/yolov8_tensorrt/weights/yolov8n-seg.trt";
    string dataset = "/media/daybeha/LTFM2/dataset/urban/urban39-pankyo/image/stereo_left";

    if(argc > 2){
      model_path = argv[1];
      dataset = argv[2];
    }
    // std::string video_path = "/home/daybeha/Documents/urban_split2.mp4";
    // std::string image_path = "/media/daybeha/LTFM2/dataset/kitti/sequences/00/image_0";
//    string dataset = "/media/daybeha/LTFM2/dataset/kitti/sequences/00";



    setRenderWindow(param);
    YOLOv8Seg yolo(param);
    std::vector<unsigned char> trt_file = yolo_utils::loadModel(model_path);
    if (trt_file.empty())
    {
        sample::gLogError << "trt_file is empty!" << std::endl;
        return -1;
    }
    if (!yolo.init(trt_file))
    {
        sample::gLogError << "initEngine() ocur errors!" << std::endl;
        return -1;
    }
    yolo.check();

    // 读取图像路径
    vector<string> vstrImageLeft, vstrImageRight;
    vector<double> vTimestamps;
//    LoadKittiImages(dataset, vstrImageLeft, vstrImageRight, vTimestamps);
    LoadUrbanIamges(dataset, vstrImageLeft);
    const int nImages = vstrImageLeft.size();


    cv::Mat cur_img;
    std::vector<cv::Mat> imgs_batch;
    imgs_batch.reserve(param.batch_size);
    sample::gLogInfo << imgs_batch.capacity() << std::endl;
    for(int ni=0; ni<nImages; ni++)
    {
        cur_img = cv::imread(vstrImageLeft[ni], cv::IMREAD_COLOR); // 读取图像
        auto start = std::chrono::steady_clock::now();
        imgs_batch.emplace_back(cur_img);
        cv::imshow("original image", cur_img);

        task(yolo, param, imgs_batch, 1);
        imgs_batch.clear();
        auto end = std::chrono::steady_clock::now();
        auto time_usage = chrono::duration_cast<chrono::milliseconds>(end - start);
        std::cout << "total time: " << time_usage.count() << "ms\n";
        cv::waitKey(0);
    }

    return  -1;
}